import json
from channels.generic.websocket import AsyncWebsocketConsumer
from .utils.spark_ai import SparkAIHandler

class AIChatConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        # 这里可以添加认证逻辑
        await self.accept()
        
    async def disconnect(self, close_code):
        pass
        
    async def receive(self, text_data):
        try:
            data = json.loads(text_data)
            
            if data['type'] == 'ping':
                await self.send(json.dumps({
                    'type': 'ping',
                    'content': 'pong'
                }))
                return
                
            if data['type'] == 'message':
                message = data['message']
                history = data.get('history', [])
                
                spark_handler = SparkAIHandler(streaming=True)
                
                async for chunk in spark_handler.generate_stream_response([*history, {
                    'role': 'user',
                    'content': message
                }]):
                    await self.send(json.dumps({
                        'type': 'message',
                        'content': chunk
                    }))
                    
        except Exception as e:
            await self.send(json.dumps({
                'type': 'error',
                'content': str(e)
            }))